package com.fwmagic.spark.ml.naivebayes

import com.fwmagic.spark.util.SparkUtils
import org.apache.spark.ml.classification.NaiveBayesModel
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
 * 朴素贝叶斯算法
 *
 * 利用朴素贝叶斯算法模型做预测
 * 1.加载算好的模型
 * 2.加载待预测的数据
 * 3.用模型预测数据结果
 * 4.输出预测结果
 */
object LoadModelAndTransform {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkUtils.getSparkSession(this.getClass.getSimpleName)

    //1.加载算好的模型数据
    val naiveBayesModel: NaiveBayesModel = NaiveBayesModel.load("data/naivebayes/model")

    //2.加载待预测的数据集
    val test: DataFrame = spark.read.option("header", true).csv("data/naivebayes/test.csv")

    test.createOrReplaceTempView("test")

    import com.fwmagic.spark.ml.utils.VectorUtils._
    spark.udf.register("arr2vec2", arr2vec2)
    //特征向量化
    val testVecs: DataFrame = spark.sql(
      """
        |
        |select
        |name,
        |arr2vec2(
        | array(job,income,age,sex)
        |) as vec
        |from test
        |
        |""".stripMargin)

    //3.用模型来预测结果
    val transformDF: DataFrame = naiveBayesModel.transform(testVecs)

    //4.输出预测结果
    transformDF.show(100, false)

    spark.close()
  }
}
